/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include <iostream>
#include <complex>
#include <vector>
#include <cmath>
#include "test_common.hpp"

// ===================== 修复比较函数 =====================
// 浮点数比较
static inline bool nearly_equal(float a, float b, float eps = 1e-5f) {
    return std::fabs(a - b) < eps;
}

// 双精度浮点数比较
static inline bool nearly_equal(double a, double b, double eps = 1e-12) {
    return std::fabs(a - b) < eps;
}

// 复数类型比较（完全重写，避免模板链接错误）
static inline bool nearly_equal(std::complex<float> a, std::complex<float> b, float eps = 1e-5f) {
    return std::fabs(a.real() - b.real()) < eps &&
           std::fabs(a.imag() - b.imag()) < eps;
}

static inline bool nearly_equal(std::complex<double> a, std::complex<double> b, double eps = 1e-12) {
    return std::fabs(a.real() - b.real()) < eps &&
           std::fabs(a.imag() - b.imag()) < eps;
}

// ===================== 向量比较函数 =====================
template <typename T>
bool compare_vectors(size_t n, const T* a, const T* b, 
                     decltype(std::declval<T>() == std::declval<T>()) eps = 
                     std::is_same<T, float>::value ? 1e-5f : 
                     std::is_same<T, double>::value ? 1e-12 :
                     std::is_same<T, std::complex<float>>::value ? 1e-5f : 1e-12) 
{
    for (size_t i = 0; i < n; ++i) {
        if (!nearly_equal(a[i], b[i], eps)) {
            std::cout << "Mismatch at position " << i << ": "
                      << "expected " << b[i] << ", got " << a[i] << std::endl;
            return false;
        }
    }
    return true;
}

// ===================== 向量打印函数 =====================
template <typename T>
void print_vector(const char* name, size_t n, const T* v) {
    std::cout << name << ": [";
    for (size_t i = 0; i < n; ++i) {
        if constexpr (std::is_same<T, std::complex<float>>::value || 
                      std::is_same<T, std::complex<double>>::value) {
            std::cout << "(" << v[i].real() << ", " << v[i].imag() << ")";
        } else {
            std::cout << v[i];
        }
        if (i < n-1) std::cout << ", ";
    }
    std::cout << "]" << std::endl;
}

// ===================== 主测试函数 =====================
TEST(SparseBLASlevel2, gather_zero) {
    int failures = 0;

    std::cout << "=== Testing OpenSpB_gather_zero (Fully Fixed Version) ===" << std::endl;

    // ===================== float (sgther_zero) 测试 =====================
    {
        const int nz = 3;
        float y_our[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
        float y_ref[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
        float x_our[3] = {0};
        float x_ref[3] = {0};
        int indx[] = {0, 2, 4}; // 0-based索引

        // 调用目标函数
        OpenSpB_sgther_zero(nz, y_our, 1, x_our, indx, OPENSPBLAS_ZERO_BASE);
        
        // 调用参考函数
        BLAS_susgz(nz, y_ref, 1, x_ref, indx, blas_zero_base);
        
        bool x_equal = compare_vectors(static_cast<size_t>(nz), x_our, x_ref);
        bool y_equal = compare_vectors(5, y_our, y_ref);
        
        if (!x_equal || !y_equal) {
            std::cout << "FAIL: sgther_zero test" << std::endl;
            print_vector("Expected x", static_cast<size_t>(nz), x_ref);
            print_vector("Actual x", static_cast<size_t>(nz), x_our);
            print_vector("Expected y", 5, y_ref);
            print_vector("Actual y", 5, y_our);
            failures++;
        } else {
            std::cout << "PASS: sgther_zero test" << std::endl;
        }
    }

    // ===================== double (dgther_zero) 测试 =====================
    {
        const int nz = 4;
        double y_our[] = {1.1, 2.2, 3.3, 4.4, 5.5};
        double y_ref[] = {1.1, 2.2, 3.3, 4.4, 5.5};
        double x_our[4] = {0};
        double x_ref[4] = {0};
        int indx[] = {1, 2, 3, 4}; // 1-based索引

        // 调用目标函数
        OpenSpB_dgther_zero(nz, y_our, 1, x_our, indx, OPENSPBLAS_ONE_BASE);
        
        // 调用参考函数
        BLAS_dusgz(nz, y_ref, 1, x_ref, indx, blas_one_base);
        
        if (!compare_vectors(static_cast<size_t>(nz), x_our, x_ref) ||
            !compare_vectors(5, y_our, y_ref)) {
            std::cout << "FAIL: dgther_zero test" << std::endl;
            failures++;
        } else {
            std::cout << "PASS: dgther_zero test" << std::endl;
        }
    }

    // ===================== complex<float> (cgther_zero) 测试 =====================
    {
        const int nz = 2;
        std::complex<float> y_our[] = {{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}};
        std::complex<float> y_ref[] = {{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}};
        std::complex<float> x_our[2] = {0};
        std::complex<float> x_ref[2] = {0};
        int indx[] = {0, 2}; // 0-based索引

        // 调用目标函数
        OpenSpB_cgther_zero(nz, y_our, 1, x_our, indx, OPENSPBLAS_ZERO_BASE);
        
        // 调用参考函数
        BLAS_cusgz(nz, y_ref, 1, x_ref, indx, blas_zero_base);
        
        if (!compare_vectors(static_cast<size_t>(nz), x_our, x_ref) ||
            !compare_vectors(3, y_our, y_ref)) {
            std::cout << "FAIL: cgther_zero test" << std::endl;
            print_vector("Expected x", static_cast<size_t>(nz), x_ref);
            print_vector("Actual x", static_cast<size_t>(nz), x_our);
            print_vector("Expected y", 3, y_ref);
            print_vector("Actual y", 3, y_our);
            failures++;
        } else {
            std::cout << "PASS: cgther_zero test" << std::endl;
        }
    }

    // ===================== 修复 complex<double> (zgther_zero) 测试 =====================
    {
        const int nz = 3;
        std::complex<double> y_our[] = {{1.1, 2.2}, {3.3, 4.4}, {5.5, 6.6}, {7.7, 8.8}};
        std::complex<double> y_ref[] = {{1.1, 2.2}, {3.3, 4.4}, {5.5, 6.6}, {7.7, 8.8}};
        std::complex<double> x_our[3] = {0};
        std::complex<double> x_ref[3] = {0};
        
        // 修复索引和步长
        int indx[] = {0, 1, 2}; // 有效的0-based索引（0-3）
        const int incy = 1;      // 标准步长

        // 调用目标函数
        OpenSpB_zgther_zero(nz, y_our, incy, x_our, indx, OPENSPBLAS_ZERO_BASE);
        
        // 调用参考函数
        BLAS_zusgz(nz, y_ref, incy, x_ref, indx, blas_zero_base);
        
        bool x_match = compare_vectors(static_cast<size_t>(nz), x_our, x_ref);
        bool y_match = compare_vectors(4, y_our, y_ref);
        
        if (!x_match || !y_match) {
            std::cout << "FAIL: zgther_zero test" << std::endl;
            failures++;
        } else {
            std::cout << "PASS: zgther_zero test" << std::endl;
        }
    }

    // ===================== 空输入测试 =====================
    {
        float y[] = {1.0f, 2.0f, 3.0f};
        float x;
        int* indx = nullptr;
        
        // 应安全处理空输入
        OpenSpB_sgther_zero(0, y, 1, &x, indx, OPENSPBLAS_ZERO_BASE);
        std::cout << "PASS: empty input test (no crash)" << std::endl;
    }

    // ===================== 测试总结 =====================
    std::cout << "\n=== Test Summary ===" << std::endl;
    if (failures == 0) {
        std::cout << "ALL TESTS PASSED!" << std::endl;
    } else {
        std::cout << failures << " TESTS FAILED!" << std::endl;
    }
}